# -*- coding: utf-8 -*-
# File : vul_scan.py
# Author: taoyahui
# Date : 2021/6/17
import sys
sys.path.append('..')
# !!!!! 漏洞扫描任务 !!!!!
import configparser
import os
from utils.path_util import get_root_dir
from utils import mysql_util, mongo_util, log_util
import toml
import logging
from results import *
from utils import vuls_utils,common_util
import json
import time
from application_factory.extensions import scheduler
from application_factory.models import InventoryCls
import sys
from fabric import Connection
import requests

# 初始化设置
logging.basicConfig(level=logging.DEBUG)
cf = configparser.ConfigParser()
cf.read(os.path.join(get_root_dir(), "config.ini"))
secs = cf.sections()
vul_dir = cf.get("vuls", "root_dir")
mongo_db = cf.get('Mongo-Database', 'db')

# 读取资产漏洞配置信息
critical_lv = int(cf.get('inventory_safe_lv', 'critical'))
high_lv = int(cf.get('inventory_safe_lv', 'high'))
medium_lv = int(cf.get('inventory_safe_lv', 'medium'))
low_lv = int(cf.get('inventory_safe_lv', 'low'))


def call_vul_scan(ip, hostname):
    result = requests.post("http://localhost:5002/vul_scan", json={"ip": ip, "hostname": hostname})

# 漏扫主入口，整合配置文件定义，漏洞扫描，漏洞分析与提取
def vul_scan_integrity(ip=None, hostname=None):
    logger = common_util.get_logger("vul_scan_integrity")
    logger.info(f"Start to Generate config file,  IP is {ip}, HOSTNAME is {hostname}")
    config_generate(ip)
    logger.info(f"Start to Excute cve_detect, IP is {ip}, HOSTNAME is {hostname}")
    vuls_utils.cve_detect()
    logger.info(f"Start to analysis vuls, IP is {ip}, HOSTNAME is {hostname}")
    analysis_vuls(hostname)
    logger.info(f"VUL Scan Finished! IP is {ip}, HOSTNAME is {hostname}")


# 通过查询数据库确定要扫描的主机
# 生成vuls扫描的配置文件config.toml
def config_generate(ip=None):
    exec_sql = 'select * from inventory where username is not null'

    if ip is not None:
        exec_sql = f"select * from inventory where username is not null and ip = '{ip}' "

    results = mysql_util.exec_sql_new(exec_sql)
    config_dict = {}
    for result in results:
        inventory = InventoryCls(result)
        ip = inventory.ip
        username = inventory.username
        hostname = inventory.hostname
        config_dict[hostname] = {'host': ip,
                           'port': '22',
                           'user': username,
                           'keyPath': '/root/.ssh/id_rsa',
                           'scanMode': ['fast', 'offline']
                           }
        # 如果操作系统时centos的话开始执行删除多余的内核(kernel)
        if 'centos' in inventory.os.lower():
            try:
                print()
                # execute_remove_dump_kernel(inventory.ip, inventory.username, inventory.password, inventory.port)
            except Exception as e:
                print(e)
    group_config = {'servers': config_dict}

    config_file_path = os.path.join(vul_dir, 'config.toml')
    logging.info(f'Config Path is {config_file_path}')
    with open(config_file_path, 'w') as f:
        r = toml.dump(group_config, f)
        # print(r)
        # logging.INFO(r)
    return success_result


def execute_remove_dump_kernel(ip, username, passwd, port):
    """
    执行删除多余内核的操作
    :param ip:
    :param username:
    :param passwd:
    :param port:
    :return:
    """
    c = Connection(f"{username}@{ip}", port=port, connect_kwargs={'password': passwd})
    try:
        log_util.logger.info(f"Start to remove dump kernel, ip is {ip}")
        c.run("yum remove -y $(rpm -qa | grep kernel | grep -v $(uname -r))")
    except Exception as e:
        log_util.logger.error(f"Excute remove dump kernel error, reason is {e}")
    finally:
        c.close()


# 漏洞分析
# 找出最新的vuls/result目录，分析里面的json
def analysis_vuls(hostname=None):
    # os.system(f"rm {vul_dir}/current")
    # vuls项目下的results目录
    results_path = f"{vul_dir}/results"

    # 如果current文件存在，则删除
    if os.path.exists(f"{results_path}/current"):
        os.system(f"rm {results_path}/current")

    # 通过时间戳生成本次搜索的id(query_id)，即本次查询的编号
    query_id = common_util.get_current_timestamp()

    # 如何hostname为None，那么就找到最近的文件夹
    if hostname is None:
        results_dirs = common_util.get_file_list(results_path)
        if results_dirs.__len__() == 0:
            return vul_detect_empty_result

        # 找到results目录下最新的文件夹
        result_latest_dir = os.path.join(results_path, results_dirs[-1])
        # 列出文件夹下所有的搜索文件
        file_names = os.listdir(result_latest_dir)

        # 遍历搜索文件
        for file_name in file_names:
            hostname = file_name[:-5]
            # print(hostname)
            json_file_path = os.path.join(result_latest_dir, file_name)
            one_host_vul_analysis(hostname=hostname,
                                  query_id=query_id,
                                  json_file_path=json_file_path)

    if hostname is not None:
        json_file_path = common_util.get_latest_file_path(results_path, hostname)
        one_host_vul_analysis(hostname=hostname,
                              query_id=query_id,
                              json_file_path=json_file_path)


def one_host_vul_analysis(hostname, query_id, json_file_path):
    inventory_ip = find_inventory_by_hostname(hostname)

    # 如果资产不存在就不往下执行了
    if inventory_ip is None:
        return

    # 删除inventory_patch对应hostname的过期数据
    remove_outdate_intentory_cve(inventory_ip)

    json_file = open(json_file_path, 'r', encoding='utf8')
    result_json = json.load(json_file)

    # 更新资产的包内容信息,通过apscheduler中的job实现(立即执行)
    save_inventory_package_job_id = f"{hostname}-{int(time.time())}"
    scheduler.add_job(func=save_inventory_package,
                      id=save_inventory_package_job_id,
                      kwargs={
                          'json_file_path': json_file_path,
                          'hostname': hostname}
                      )
    # save_inventory_package(result_json, hostname)

    # 获取scannedCves信息
    if 'scannedCves' in result_json.keys():


        patch_info_list = []
        for cve_item in result_json['scannedCves'].values():
            cve_id = cve_item['cveID']

            # 生成补丁信息内容，注意：其中不包含影响包的内容
            # 补丁对应包修复包的信息将分别存在每个资产对应的补丁信息中
            # patch表中只存大致的信息，比如名称，描述，危险等级等信息
            patch_info = {'cve_id': cve_id, 'patch_no': None}

            # 判断是否对应的解决方案
            if 'distroAdvisories' in cve_item.keys():
                patch_info['patch_no'] = cve_item['distroAdvisories'][0]['advisoryID']
                patch_info['abs'] = cve_item['distroAdvisories'][0]['description']
                patch_info['publishing_time'] = input_date_format(cve_item['distroAdvisories'][0]['issued'])  # 发布时间

            if 'cveContents' in cve_item.keys():
                cve_contents = cve_item['cveContents']
                if 'nvd' in cve_contents.keys():
                    patch_info['cve_summary'] = cve_contents['nvd']['summary']
                    patch_info['cve_cvss3_severity'] = cve_contents['nvd']['cvss3Severity']
                    patch_info['cve_cvss3_score'] = cve_contents['nvd']['cvss3Score']
                    patch_info['cve_publish_time'] = input_date_format(cve_contents['nvd']['published'])
                    patch_info['cve_modify_time'] = input_date_format(cve_contents['nvd']['lastModified'])
                    patch_info['cve_source_link'] = cve_contents['nvd']['sourceLink']
                    patch_info['danger_lv'] = set_danger_level(patch_info['cve_cvss3_severity'])

                if 'redhat' in cve_contents.keys():
                    patch_info['cve_title'] = cve_contents['redhat']['title']
                    patch_info['publisher'] = 'redhat'

            # 绑定cve和inventory资产
            bind_patch_inventory(inventory_ip, cve_id, patch_info['patch_no'], query_id,
                                 json.dumps(cve_item['affectedPackages']))
            patch_info['cve_affect_packages'] = json.dumps(cve_item['affectedPackages'])
            patch_info['hostname'] = hostname
            patch_info_list.append(patch_info)
            # logging.log(logging.DEBUG, f"Patch {patch_info['cve_id']} saved to list")
            # print(f"Patch {patch_info['cve_id']} saved to list")

        # 将资产的补丁信息保存至mongodb
        print("开始存储数据至Mongodb")
        save_inventory_patch_to_mongo(patch_info_list)

        update_inventory_cve_sql = f"Update inventory set unpatch_num = {len(patch_info_list)} where hostname = '{hostname}'"
        mysql_util.exec_sql_new(update_inventory_cve_sql)

        # 根据搜索结果更新资产表中对应资产的"未打补丁的信息"
        # pack_length = len(patch_info_list)
        # update_inventory_cve_sql = f"Update inventory set unpatch_num = {pack_length} where hostname = '{hostname}'"
        # mysql_util.exec_sql_new(update_inventory_cve_sql)
        update_inventory_unpatch_num(inventory_ip=inventory_ip)


        # 根据规则修改资产的安全等级
        update_inventory_safe_lv(inventory_ip)


# 在mongodb中保存资产和补丁信息
def save_inventory_patch_to_mongo(patch_info_list):
    if len(patch_info_list) > 0:
        # logging.log(logging.INFO, f"common_util.gen_error_msg(sys._getframe().f_code.co_name) Start to save patch, hostname is {hostname}")
        res_id = mongo_util.insert_many_document(db_name=mongo_db, collection_name='inventory_patch', doc= patch_info_list)


# 设置危险等级
def set_danger_level(severity):
    if severity == 'CRITICAL':
        return 3
    elif severity == 'HIGH':
        return 2
    elif severity == 'MEDIUM':
        return 1
    else:
        return 0


# !!!!! 此操作可能引起数据误报（删除inventory_patch数据表中的多余数据） !!!!!
# 只保留现有数据库中最新query_id的数据
def remove_outdate_intentory_cve(inventory_ip):
    query_id_group_sql = f"select query_id from inventory_patch where host = '{inventory_ip}' and query_id is not null GROUP BY query_id order by query_id desc"
    query_res_list = mysql_util.exec_sql_new(query_id_group_sql)

    if query_res_list is None:
        return 0

    # 找出现有数据库中的所有
    for res in query_res_list[1:]:
        query_id = res[0]
        delete_res_sql = f"DELETE FROM inventory_patch where query_id = {query_id} and host = '{inventory_ip}'"
        mysql_util.exec_sql_new(delete_res_sql)
    return 1


# 绑定漏洞和资产关系
def bind_patch_inventory(hostname, cve_id, patch_no, query_id, affect_packages):
    conn = mysql_util.get_conn()
    cursor = mysql_util.get_cursor(conn)
    insert_dict = {
        'host': hostname,
        'cve_id': cve_id,
        'patch_no': patch_no,
        'utime': common_util.time_str_format(),
        'query_id':query_id,
        'affect_packages': affect_packages
    }
    insert_sql = mysql_util.generate_insert_sql(insert_dict, 'inventory_patch')
    mysql_util.exec_sql_params(insert_sql, tuple(insert_dict.values()))
    mysql_util.release(conn, cursor)


# 通过主机名找资产ip等信息
def find_inventory_by_hostname(hostname):
    sql = f"select ip from inventory where hostname = '{hostname}'"
    result = mysql_util.exec_sql_new(sql)

    if result is None:
        logging.info(f"Inventory {hostname} not exist!")
        return None
    ip = result[0][0]
    return ip


# 将格式为'%Y-%m-%dT%H:%M:%SZ'的字符串 转换为格式为'%Y-%m-%d %H:%M:%S' 的字符串
def input_date_format(time_str):
    format_time = time.strptime(time_str, '%Y-%m-%dT%H:%M:%SZ')
    new_format_time = time.strftime('%Y-%m-%d %H:%M:%S', format_time)
    return new_format_time


# 保存资产的包信息
# 表名为: 资产的hostname
def save_inventory_package(json_file_path, hostname):
    """
    存储资产对应的软件包
    :param json_file_path:
    :param hostname:
    :return:
    """
    json_file = open(json_file_path, 'r', encoding='utf8')
    result_json = json.load(json_file)
    packages = result_json['packages']
    # 将数据存入mongodb
    input_docs = []
    mongo_util.truncate_collection(db_name=mongo_db, collection_name=hostname)
    for pac_name in packages.keys():
        input_docs.append(packages[pac_name])
    mongo_util.insert_many_document(mongo_db, hostname, input_docs)

    # 将package包数据数据存入 mysql数据库
    # 1. 首先清空inventory_package相关hostname的数据
    delete_sql = f"Delete From inventory_package where hostname = '{hostname}'"
    mysql_util.exec_sql_new(delete_sql)
    for index in packages:
        pac_dict = {
            'hostname': hostname,
            'name': packages[index]['name'],
            'version': packages[index]['version'],
            'release_no': packages[index]['release'],
            'new_version': packages[index]['newVersion'],
            'new_release_no': packages[index]['newRelease'],
            'arch': packages[index]['arch'],
            'repository': packages[index]['repository']
        }
        pac_insert_sql = mysql_util.generate_insert_sql(pac_dict, 'inventory_package')
        mysql_util.exec_sql_params(pac_insert_sql, tuple(pac_dict.values()))
    return 1


def update_inventory_safe_lv(inventory_ip):
    """
    更新资产的安全等级0: 安全 1:低危 2:中危 3:高危
    :param inventory_ip: 资产ip
    :return:
    """
    count_patch_sql = f"select count(DISTINCT(a.cve_id)) as patch_c, query_id  from inventory_patch as a  " \
                      f"left join patch as b on a.cve_id = b.cve_id " \
                      f"where a.host = '{inventory_ip}' and (b.cve_cvss3_severity = 'HIGH' or b.cve_cvss3_severity = 'CRITICAL') " \
                      f"GROUP BY a.query_id"
    res_list = mysql_util.exec_sql_new(count_patch_sql)
    if len(res_list) == 0:
        return
    patch_count = res_list[-1][0]
    safe_lv = 0
    if patch_count > critical_lv:
        safe_lv = 3
    elif patch_count > high_lv:
        safe_lv = 2
    elif patch_count > medium_lv:
        safe_lv = 1
    if safe_lv > 0:
        update_sql = f"Update inventory set safe_lv = {safe_lv} where ip = '{inventory_ip}'"
        mysql_util.exec_sql_new(update_sql)


def update_inventory_unpatch_num(inventory_ip):
    """
    更新资产未打补丁的数量(包括这次漏扫和上次漏扫)
    :param inventory_ip:  主机的ip
    :return:
    """
    count_patch_sql = f"select count(DISTINCT(a.cve_id)) as patch_c, query_id  from inventory_patch as a  " \
                      f"left join patch as b on a.cve_id = b.cve_id " \
                      f"where a.host = '{inventory_ip}'" \
                      f"GROUP BY a.query_id"

    res_list = mysql_util.exec_sql_new(count_patch_sql)
    if len(res_list) > 1:
        this_unpatch_num = res_list[-1][0]
        last_unpatch_num = res_list[-2][0]
        update_sql = f"Update inventory set unpatch_num = '{this_unpatch_num}', last_unpatch_num = '{last_unpatch_num}' " \
                     f"where ip = '{inventory_ip}'"
        mysql_util.exec_sql_new(update_sql)
    elif len(res_list) > 0:
        this_unpatch_num = res_list[-1][0]
        update_sql = f"Update inventory set unpatch_num = '{this_unpatch_num}' " \
                     f"where ip = '{inventory_ip}'"
        mysql_util.exec_sql_new(update_sql)


if __name__ == '__main__':
    # pass
    # import datetime
    # print(datetime.datetime.now())
    # vul_scan_integrity()
    # print(datetime.datetime.now())
    # config_generate('192.168.50.156')
    # vuls_utils.cve_detect()
    # analysis_vuls()
    # remove_outdate_intentory_cve('123')
    # bind_cve_inventory('test', 'CVE-2020-12770', 'RHSA-928-098', common_util.get_current_timestamp())
    # print(input_date_format('2021-06-07T00:00:00Z'))
    # update_inventory_safe_lv('192.168.50.156')
    execute_remove_dump_kernel('192.168.50.125', 'root', '123456', 22)

